import io

import matplotlib.pyplot as plt
import pandas as pd
from fastapi import HTTPException


class DataAnalysis:
    def __init__(self, file_path: str):
        self.file_path = file_path
        self.df = None
        self.load_data()

    # 加载、预处理数据集
    def load_data(self):
        try:
            # 加载数据
            excel_file = pd.ExcelFile(self.file_path)
            self.df = excel_file.parse("Sheet1")
            # 数据清洗
            self.df = self.df.drop_duplicates()
            self.df["Patient Sex"] = self.df["Patient Sex"].map({"Male": 1, "Female": 0})
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"数据加载失败: {str(e)}")

    # 患者基本特征
    def get_patient_features(self) ->dict:
        age = self.df["Patient Age"].describe().round(2).to_dict()
        sex = self.df["Patient Sex"].value_counts().to_dict()
        sex = {"Male": sex.get(1,0), "Female": sex.get(0,0)}
        return {"age": age, "sex": sex}

    # 疾病分布情况
    def get_disease_distribution(self) -> dict:
        disease_columns = ['N','D','G','C','A','H','M','O']
        disease_labels = {'N': '正常', 'D': '糖尿病', 'G': '青光眼', 'C': '白内障', 'A': 'AMD', 'H': '高血压', 'M': '近视', 'O': '其他疾病/异常'}
        disease_proportions = self.df[disease_columns].mean().round(4)*100
        disease_data = {
            disease_labels.get(col, col): value
            for col, value in disease_proportions.items()
        }
        return disease_data

    # 分析疾病和年龄的相关性
    def age_disease_correlation(self) -> dict:
        disease_columns = ['N','D','G','C','A','H','M','O']
        disease_labels = {'N': '正常', 'D': '糖尿病', 'G': '青光眼', 'C': '白内障', 'A': 'AMD', 'H': '高血压', 'M': '近视', 'O': '其他疾病/异常'}
        correlations = self.df[['Patient Age']+disease_columns].corr()['Patient Age'][1:].round(4).to_dict()
        return {disease_labels.get(col, col): value for col, value in correlations.items()}

    # 生成年龄分布表
    def generate_age_distribution_chart(self) -> str:
        try:
            plt.figure(figsize=(10, 6))
            plt.hist(self.df['Patient Age'], bins=20)
            plt.title("患者年龄分布")
            plt.xlabel("年龄")
            plt.ylabel("数量")
            buffer = io.BytesIO ()
            plt.savefig(buffer, format='png')
            buffer.seek(0)
            plt.close()
            return buffer
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"生成年龄分布图表失败: {str(e)}")

    # 疾病分布图
    def generate_disease_distribution_chart(self) -> str:
        try:
            disease_data = self.get_disease_distribution()
            plt.figure(figsize=(10, 8))
            plt.pie(
                disease_data.values(),
                labels=disease_data.keys(),
                autopct='%1.1f%%',
                startangle=90
            )
            plt.axis('equal')
            plt.title("疾病分布")
            buffer = io.BytesIO ()
            plt.savefig(buffer, format='png')
            buffer.seek(0)
            plt.close()
            return buffer
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"生成疾病分布图表失败: {str(e)}")
